'''
无条件图像生成 推理

'''
import time

from diffusers import DDPMScheduler, UNet2DModel
from tqdm import tqdm

# load model, scheduler
# model_id = "stabilityai/stable-diffusion-2-1"
# model_id = "google/ddpm-cat-256"
# model_id = "ddpm-butterflies-128"
model_id = "zhangxiancai/strap2-unet-model"
scheduler = DDPMScheduler.from_pretrained(model_id, subfolder='scheduler')
model = UNet2DModel.from_pretrained(model_id, subfolder='unet',use_safetensors=True).to("cuda")

scheduler.set_timesteps(25)

import torch
sample_size = model.config.sample_size
print(f'img size {sample_size}')
generator = torch.cuda.manual_seed(100) #
noise = torch.randn((1, 3, sample_size, sample_size), device="cuda",generator=generator)

# 迭代去噪
input = noise
# for t in scheduler.timesteps:
for t in tqdm(scheduler.timesteps):
    with torch.no_grad():
        noisy_residual = model(input, t).sample
    previous_noisy_sample = scheduler.step(noisy_residual, t, input).prev_sample
    input = previous_noisy_sample



from PIL import Image
image = (input / 2 + 0.5).clamp(0, 1).squeeze()
image = (image.permute(1, 2, 0) * 255).round().to(torch.uint8).cpu().numpy()
image = Image.fromarray(image)
image.save(f"demo_{time.time()}.png")